import sensor, image, lcd, time
import KPU as kpu
import gc, sys
from machine import UART
from Maix import GPIO
from fpioa_manager import fm

board_cube = 1
UART_TX = 24  # 6 24
UART_RX = 25 # 7 25
api_run = 0
anchors = (0.1606, 0.3562, 0.4712, 0.9568, 0.9877, 1.9108, 1.8761, 3.5310, 3.4423, 5.6823)
model_addr="/sd/models/m5/m.smodel"
class_IDs = ['no_mask', 'mask']
data_send= ''

color_R = (255, 0, 0)
color_G = (0, 255, 0)
color_B = (0, 0, 255)

def drawConfidenceText(image, rol, classid, value):
    text = ""
    _confidence = int(value * 100)

    if classid == 1:
        text = 'mask: ' + str(_confidence) + '%'
        color_text=color_G
    else:
        text = 'no_mask: ' + str(_confidence) + '%'
        color_text=color_R
    image.draw_string(rol[0]+48, rol[1], text, color=color_text, scale=2.5)

def lcd_show_except(e):
    import uio
    err_str = uio.StringIO()
    sys.print_exception(e, err_str)
    err_str = err_str.getvalue()
    img = image.Image(size=(224,224))
    img.draw_string(0, 10, err_str, scale=1, color=(0xff,0x00,0x00))
    lcd.display(img)

def loop(uart_A=None,key_enter=1,key_back=1,key_next=1,kpu=None,task=None):
    try:
        global data_send
        img = sensor.snapshot()
        if board_cube:
            img = img.rotation_corr(z_rotation=90)
            img.pix_to_ai()
        t = time.ticks_ms()
        code = kpu.run_yolo2(task, img)
        t = time.ticks_ms() - t
        if code:
            totalRes = len(code)
            #print("had face ")
            for item in code:
                confidence = float(item.value())
                itemROL = item.rect()
                classID = int(item.classid())

                if confidence < 0.52:
                    _ = img.draw_rectangle(itemROL, color=color_B, tickness=5)
                    continue

                if classID == 1 and confidence > 0.65:
                    _ = img.draw_rectangle(itemROL, color_G, tickness=5)
                    if totalRes == 1:
                        drawConfidenceText(img, (0, 0), 1, confidence)
                        print("mask")
                        if data_send != 'mask':
                            uart_A.write('mask')
                            data_send = 'mask'
                else:
                    _ = img.draw_rectangle(itemROL, color=color_R, tickness=5)
                    if totalRes == 1:
                        drawConfidenceText(img, (0, 0), 0, confidence)
                        print("no_mask")
                        if data_send != 'no_mask':
                            uart_A.write('no_mask')
                            data_send = 'no_mask'

        else:
            if data_send != 'None':
                uart_A.write('None')
                data_send = 'None'
        img.draw_string(48, 210, "Mask t:%dms" %(t), scale=2)
        lcd.display(img)
        time.sleep_ms(100)
    except Exception as e:
        raise e
    finally:
        del img

def work():
    try:
        uart_A = UART(UART.UART1, 115200, 8, 1, 0, timeout=1000, read_buf_len=4096)
        # reg
        fm.register(16, fm.fpioa.GPIOHS0) # 16 NEXT
        fm.register(11, fm.fpioa.GPIOHS1) #11 BACK
        fm.register(10, fm.fpioa.GPIOHS2) #10 ENTER
        key_next = GPIO(GPIO.GPIOHS0, GPIO.PULL_UP) # 0
        key_back = GPIO(GPIO.GPIOHS1, GPIO.PULL_UP) # 0
        key_enter = GPIO(GPIO.GPIOHS2, GPIO.PULL_UP) # 0

        fm.register(UART_TX, fm.fpioa.UART1_TX, force=True) # 6 24
        fm.register(UART_RX, fm.fpioa.UART1_RX, force=True) # 7 25
        time.sleep_ms(100) # wait uart ready

        uart_A.write(b'hello world')

        task = None
        task = kpu.load(model_addr)
        #init yolo2 ,0.5（50%）
        kpu.init_yolo2(task, 0.5, 0.3, 5, anchors) # threshold:[0,1], nms_value: [0, 1]

        while(True):
            loop(uart_A,key_enter,key_back,key_next,kpu,task)
    except Exception as e:
        sys.print_exception(e)
        lcd_show_except(e)
    finally:
        if not task is None:
            kpu.deinit(task)
        fm.unregister(16)
        fm.unregister(11)
        fm.unregister(10)
        uart_A.deinit()
        del uart_A
        gc.collect()

def main(lcd_rotation=0, sensor_hmirror=False, sensor_vflip=False):
    try:
        sensor.reset()
    except Exception as e:
        raise Exception("sensor reset fail, please check hardware connection, or hardware damaged! err: {}".format(e))

    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.QVGA)
    #sensor.skip_frames()
    #sensor.set_windowing((224, 224))
    if board_cube == 1:
        sensor.set_vflip(True)
        sensor.set_hmirror(True)
        lcd.init(type=2)
        lcd.rotation(2)
    else:
        lcd.init(type=1)
        lcd.rotation(lcd_rotation)
        sensor.set_hmirror(sensor_hmirror)
        sensor.set_vflip(sensor_vflip)
        lcd.rotation(lcd_rotation)
    sensor.run(1)
    lcd.clear(lcd.WHITE)
    work()

# api
def run(uart_A=None,key_enter=1,key_back=1,key_next=1):
    print("start Mask recognition")
    try:
        global api_run
        api_run = 1
        task = None
        task = kpu.load(model_addr)
        #init yolo2 ,0.5（50%）
        kpu.init_yolo2(task, 0.5, 0.3, 5, anchors)
        while(api_run):
            loop(uart_A,key_enter,key_back,key_next,kpu,task)
    except Exception as e:
        sys.print_exception(e)
        lcd_show_except(e)
    finally:
        if not task is None:
            kpu.deinit(task)
        gc.collect()

if __name__ == "__main__":
    print("start Mask recognition")
    #model_addr=0x500000
    main()
